# coding: utf-8
# create by tongshiwei on 2018/11/10
from RecAgent.Agent import Agent

import random
import numpy as np
from collections import Counter


class MCAgent(Agent):
    def __init__(self, graph, kt, mc_times, model_dir, path_warm_up=True, window_size=None, visit_threshold=1,
                 allow_shortcut=True):

        super(MCAgent, self).__init__(graph)

        self.kt = kt
        self.mc_times = mc_times
        self.max_steps = 0
        self.model_dir = model_dir

        self._state = None
        self._target = None
        self.path = None
        self.pre_path = None
        self._legal_candidates = None

        self.window_size = window_size
        self.visit_threshold = visit_threshold
        self.allow_shortcut = allow_shortcut
        self.path_warm_up = path_warm_up

        self._cnt = 0

    @property
    def kt_state(self):
        return self._state[1]

    @property
    def mastery(self):
        return self._state[0]

    @property
    def state(self):
        return self.mastery

    def begin_episode(self, exercises_record, target, *args, **kwargs):
        mastery, _state = self.kt(exercises_record)
        self._state = (mastery[-1], _state)
        self._target = target if target is not None else list(self.graph.nodes)
        self.path = []
        self.pre_path = [e for e, c in exercises_record]
        self._legal_candidates = set(target)

    def test_correct(self, exercise, mastery):
        return 1 if 0.5 <= mastery[exercise] else 0

    def test(self, exercises, score_type=False, mastery=None):
        mastery = mastery if mastery is not None else self.mastery
        if score_type:
            return [(exercise, mastery[exercise]) for exercise in exercises]
        return [(exercise, self.test_correct(exercise, mastery)) for exercise in exercises]

    def test_score(self, exercises, score_type=False, mastery=None):
        return sum([s for _, s in self.test(exercises, score_type, mastery)])

    def search(self):
        initial_state = self._state

        delta_base = len(self._target) - self.test_score(self._target, mastery=self.mastery)

        assert self.max_steps > 0

        candidates = []
        print()
        for i in range(self.mc_times):
            self._state = initial_state
            rec_path = []
            for j in range(1, self.max_steps + 1):
                _exercise, _ = self.get_candidate_actions(rec_path)
                exercise = random.choice(_exercise)
                mastery, state = self.kt([(exercise, 1 if self.mastery[exercise] >= 0.5 else 0)], self.kt_state)
                self._state = (mastery[-1], state)
                rec_path.append(exercise)
            delta = len(self._target) - self.test_score(self._target, mastery=self.mastery)

            _alpha = delta / len(self._target)
            _beta = _alpha / len(rec_path)
            _gamma = delta / delta_base if delta_base > 0 else 0
            _theta = _gamma / len(rec_path)

            candidates.append([np.sign(_gamma) * np.abs(_alpha * _beta * _gamma * _theta), rec_path])
        print()
        _, rec_path = max(candidates, key=lambda x: x[0])
        return rec_path

    def step(self, *args, **kwargs):
        self.max_steps += 1
        return -1, None

    def state_transform(self, *args, **kwargs):
        pass

    def end_episode(self):
        rec_path = self.search()
        self._state = None
        self._target = None
        self.path = None
        self.max_steps = 0
        self.pre_path = None
        self._legal_candidates = None
        return rec_path

    def get_candidate_actions(self, path=None):
        path = self.path if path is None else path
        if path:
            pnode = path[-1]
        elif self.path_warm_up is True:
            pnode = self.pre_path[-1]
        else:
            pnode = None

        visit_nodes = Counter(self.path[-self.window_size:]) if self.window_size is not None else None
        return self.graph(
            self.mastery,
            pnode=pnode,
            visit_nodes=visit_nodes,
            visit_threshold=self.visit_threshold,
            allow_shortcut=self.allow_shortcut,
            target=self._target,
            legal_candidates=self._legal_candidates,
        )


def main(dataset, ctx, mc_times, max_steps, max_num, score_type, allow_shortcut=True, test_pre='', disable=False):
    import os
    from Graph import Graph
    from Environment import load_evaluator, load_dkt
    from mxnet import gpu
    from longling import wf_open

    model_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../data/%s/" % dataset)) + os.sep

    ctx = gpu(ctx)
    evaluator = load_evaluator(dataset, ctx, score_type,
                               logger=config_logging(logger="evaluator", console_log_level=LogLevel.INFO))
    kt = load_dkt(dataset, ctx)

    graph = Graph(dataset=dataset, disable=disable)
    agent = MCAgent(
        graph=graph,
        kt=kt,
        mc_times=mc_times,
        model_dir=model_dir,
        path_warm_up=True,
        allow_shortcut=allow_shortcut,
    )
    method = 'MC'
    if graph_disable is True:
        method += '-ng'
    else:
        if allow_shortcut is False:
            method += '-nsc'
    if score_type is True:
        method += '-score'
    prefix = model_dir + "Baseline/%s/PL-%s/MC%s-%s-%s" % (method, max_steps, test_pre, max_num, mc_times)
    with wf_open(prefix + "_result.log") as wf, wf_open(prefix + "_path.log") as pwf:
        print(evaluator.eval(
            agent=agent,
            max_steps=max_steps,
            epsilon=0,
            alternative_mode=False,
            log_f=pwf,
            max_num=max_num,
            delay_estimation=True,
        ), file=wf)

    return evaluator


if __name__ == '__main__':
    import sys

    assert len(sys.argv) == 6

    dataset = sys.argv[1]
    ctx = int(sys.argv[2])
    mc_times = int(sys.argv[3])
    max_steps = int(sys.argv[4])
    test_pre = sys.argv[5]

    # dataset = "irt"
    # ctx = 0
    # mc_times = 1
    # max_steps = 20
    # test_pre = "0"

    max_num = 1000
    score_type = False
    allow_shortcut = False
    graph_disable = True

    evaluator = main(
        dataset=dataset,
        ctx=ctx,
        mc_times=mc_times,
        max_steps=max_steps,
        max_num=max_num,
        score_type=score_type,
        allow_shortcut=allow_shortcut,
        test_pre=test_pre,
        disable=graph_disable,
    )
